package queryfrontend

import (
	"context"

	"connectrpc.com/connect"
	"github.com/go-kit/log/level"
	"github.com/grafana/dskit/tenant"

	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
	"github.com/grafana/pyroscope/pkg/featureflags"
	"github.com/grafana/pyroscope/pkg/validation"
)

func (q *QueryFrontend) LabelNames(
	ctx context.Context,
	c *connect.Request[typesv1.LabelNamesRequest],
) (*connect.Response[typesv1.LabelNamesResponse], error) {
	tenantIDs, err := tenant.TenantIDs(ctx)
	if err != nil {
		return nil, connect.NewError(connect.CodeInvalidArgument, err)
	}
	empty, err := validation.SanitizeTimeRange(q.limits, tenantIDs, &c.Msg.Start, &c.Msg.End)
	if err != nil {
		return nil, connect.NewError(connect.CodeInvalidArgument, err)
	}
	if empty {
		return connect.NewResponse(&typesv1.LabelNamesResponse{}), nil
	}

	labelSelector, err := buildLabelSelectorFromMatchers(c.Msg.Matchers)
	if err != nil {
		return nil, connect.NewError(connect.CodeInvalidArgument, err)
	}
	report, err := q.querySingle(ctx, &queryv1.QueryRequest{
		StartTime:     c.Msg.Start,
		EndTime:       c.Msg.End,
		LabelSelector: labelSelector,
		Query: []*queryv1.Query{{
			QueryType:  queryv1.QueryType_QUERY_LABEL_NAMES,
			LabelNames: &queryv1.LabelNamesQuery{},
		}},
	})
	if err != nil {
		return nil, err
	}
	if report == nil {
		return connect.NewResponse(&typesv1.LabelNamesResponse{}), nil
	}

	labelNames := report.LabelNames.LabelNames
	if capabilities, ok := featureflags.GetClientCapabilities(ctx); !ok || !capabilities.AllowUtf8LabelNames {
		// Filter out label names not passing legacy validation if utf8 label names not enabled
		filteredLabelNames := make([]string, 0, len(labelNames))
		for _, labelName := range labelNames {
			if _, _, ok := validation.SanitizeLegacyLabelName(labelName); !ok {
				level.Debug(q.logger).Log("msg", "filtering out label", "label_name", labelName)
				continue
			}
			filteredLabelNames = append(filteredLabelNames, labelName)
		}
		labelNames = filteredLabelNames
	}

	return connect.NewResponse(&typesv1.LabelNamesResponse{Names: labelNames}), nil
}
